library(dplyr)
library(tidyr)
library(ggplot2)
library(ggalluvial)

if (!file.exists("output")) {
  dir.create("output")
}

# load data

source("HNC_metadata_tidy.R")
source("data_handling.R")

cvCOSMIC <- read.csv("Input_signature_attributions/Input_CN_COSMIC_attributions.csv", stringsAsFactors = T) %>% rename(donor_id = X)

# CN clusters generated in Figure_7c.R
additional_df <- list(cvCOSMIC, CN_clusters <- read.csv("output/CN_clusters.csv"))

for (df in additional_df) {
  data <- data %>% left_join(df, by = "donor_id")
}

# Dichotomize the signatures (Presence (signatures>0)="1",absence (signatures==0)="0")
# If signature is present in more than 75% cases, dichotomize by the median
data <- dichotomizeSigs(metadata = data, sigsn = cvCOSMIC, sigcutoff = 2)

# Calculate frequencies by CN cluster
cn_freqs <- data %>%
  filter(!is.na(HNC_clusters)) %>%
  mutate(
    riskfactor = case_when(
      hpv_pos_opc == "Positive" ~ "HPV",
      tob_alc == "None" ~ "Unknown",
      tob_alc %in% c("Tobacco", "Tobacco + alcohol") ~ "Tobacco \n tobacco + alcohol",
      tob_alc == "Alcohol" ~ "Alcohol"
    ),
    riskfactor = factor(riskfactor, levels = c("Unknown", "HPV", "Tobacco \n tobacco + alcohol", "Alcohol")),
    HNC_clusters = factor(HNC_clusters,
      levels = c("Cluster D1", "Cluster D2", "Cluster P1", "Cluster P2"),
      labels = c("D1 \n Diploid CIN", "D2 \n CN neutral", "P1 \n 2x WGD", "P2 \n WGD CIN")
    )
  ) %>%
  group_by(HNC_clusters, riskfactor) %>%
  tally() %>%
  mutate(perc = prop.table(n) * 100, .after = "n")

ggplot(data = cn_freqs, aes(axis1 = riskfactor, axis2 = HNC_clusters, y = perc, fill = HNC_clusters)) +
  geom_alluvium(width = 0.5) +
  geom_stratum(width = 0.5) +
  geom_text(stat = "stratum", aes(label = after_stat(stratum)), size = 4) +
  scale_x_discrete(limits = c("riskfactor", "HNC_clusters"), labels = c("Risk Factor", "CN cluster"), expand = c(.2, .05)) +
  scale_fill_manual(values = c("#2c7fb8", "#41B6C4", "#A1DAB4", "#edf8b1"), na.value = "white") +
  ylab("") +
  theme_minimal() +
  theme(
    legend.position = "none",
    panel.grid.major = element_blank(),
    axis.text = element_blank()
  )

ggsave(paste0("output/Figure_7d_", Sys.Date(), ".pdf"), device = "pdf", width = 5.7, height = 5.1)
